# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import os

from absl import app, flags

import pandas
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

FLAGS = flags.FLAGS


def plot_price_diff(dfs, csv_root, product):
  plt.figure(figsize=(20, 10))
  plt.title('%s true bid ask price' % product)
  bid0, = plt.plot(dfs[0][0]['true_bid'], color='blue')
  ask0, = plt.plot(dfs[0][0]['true_ask'], color='red')

  bid1, = plt.plot(dfs[1][0]['true_bid'], color='green')
  ask1, = plt.plot(dfs[1][0]['true_ask'], color='yellow')

  bid2, = plt.plot(dfs[2][0]['true_bid'], color='brown')
  ask2, = plt.plot(dfs[2][0]['true_ask'], color='orange')

  bid3, = plt.plot(dfs[3][0]['true_bid'], color='black')
  ask3, = plt.plot(dfs[3][0]['true_ask'], color='pink')

  bid4, = plt.plot(dfs[4][0]['true_bid'], color='aliceblue')
  ask4, = plt.plot(dfs[4][0]['true_ask'], color='purple')
  plt.legend([bid0, ask0, bid1, ask1, bid2, ask2, bid3, ask3, bid4, ask4],
             [
                 '%s_true_bid' % dfs[0][1],
                 '%s_true_ask' % dfs[0][1],
                 '%s_true_bid' % dfs[1][1],
                 '%s_true_ask' % dfs[1][1],
                 '%s_true_bid' % dfs[2][1],
                 '%s_true_ask' % dfs[2][1],
                 '%s_true_bid' % dfs[3][1],
                 '%s_true_ask' % dfs[3][1],
                 '%s_true_bid' % dfs[4][1],
                 '%s_true_ask' % dfs[4][1]
             ])
  plt.ylabel('price USD/USDT')
  plt.savefig(os.path.join(csv_root, '%s_price_plot.png' % product))


def read_csv_into_df(csv_root, csv):
  csv_dir = os.path.join(csv_root, csv)
  df = pandas.read_csv(csv_dir, sep=',', header=0)
  df['datetime'] = pandas.to_datetime(df['timestamp'], unit='ns')
  df = df.set_index('datetime')
  return df


def main(argv):
  csv_root = FLAGS.csv_root
  assert csv_root, '--csv_root must be specified.'

  products = ['BCH-USD', 'BTC-USD', 'EOS-USD']  #'ETC-USD', 'ETH-USD', 'XRP-USD'
  for product in products:
    dfs = []
    kraken_price_df = read_csv_into_df(csv_root, '%s.%s.csv' % (product, 'Kraken'))
    dfs.append((kraken_price_df, 'Kraken'))
    for exchange in ['Binance', 'Bitfinex', 'Huobi', 'Okex']:
      price_df = read_csv_into_df(csv_root, '%sT.%s.csv' % (product, exchange))
      dfs.append((price_df, exchange))
    plot_price_diff(dfs, csv_root, product)

  return 0


if __name__ == '__main__':
  flags.DEFINE_string('csv_root', None, 'Input csv files root directory.')

  app.run(main)
